from flask import Flask, render_template, request, jsonify, Response
import requests
import json
import os

app = Flask(__name__)

# Ollama 的 API 配置
OLLAMA_HOST = os.getenv('OLLAMA_HOST', 'localhost')
OLLAMA_PORT = os.getenv('OLLAMA_PORT', '11434')
OLLAMA_MODEL = os.getenv('OLLAMA_MODEL', 'deepseek-r1:1.5b')
OLLAMA_URL = f"http://{OLLAMA_HOST}:{OLLAMA_PORT}/api/generate"

def ask_deepseek(question, stream=False):
    """调用本地 Ollama 的 DeepSeek 模型"""
    data = {
        "model": OLLAMA_MODEL,
        "prompt": question,
        "stream": stream
    }
    
    try:
        if stream:
            response = requests.post(OLLAMA_URL, json=data, timeout=120, stream=True)
            response.raise_for_status()
            return response
        else:
            response = requests.post(OLLAMA_URL, json=data, timeout=120)
            response.raise_for_status()
            result = response.json()
            answer = result.get("response", "抱歉，没有获取到回答")
            # 移除 <think> 标签及其内容
            import re
            answer = re.sub(r'<think>.*?</think>', '', answer, flags=re.DOTALL)
            return answer.strip()
    except requests.exceptions.ConnectionError:
        return "错误：无法连接到 OLLAMA 服务，请确保 OLLAMA 正在运行"
    except requests.exceptions.Timeout:
        return "错误：请求超时（超过2分钟），模型可能需要更长时间处理，请稍后重试"
    except requests.exceptions.RequestException as e:
        return f"请求错误：{str(e)}"
    except Exception as e:
        return f"未知错误：{str(e)}"

@app.route('/')
def index():
    return render_template('index.html')

@app.route('/health')
def health():
    """检查 OLLAMA 服务状态"""
    try:
        response = requests.get(f"http://{OLLAMA_HOST}:{OLLAMA_PORT}/api/tags", timeout=5)
        response.raise_for_status()
        models = response.json().get('models', [])
        return jsonify({
            'status': 'healthy',
            'ollama_url': OLLAMA_URL,
            'current_model': OLLAMA_MODEL,
            'available_models': [model['name'] for model in models]
        })
    except Exception as e:
        return jsonify({
            'status': 'unhealthy',
            'error': str(e),
            'ollama_url': OLLAMA_URL
        }), 500

@app.route('/ask', methods=['POST'])
def ask():
    question = request.json.get('question', '')
    if not question:
        return jsonify({'error': '问题不能为空'})
    
    answer = ask_deepseek(question)
    return jsonify({'question': question, 'answer': answer})

@app.route('/ask_stream', methods=['POST'])
def ask_stream():
    """流式响应端点"""
    question = request.json.get('question', '')
    if not question:
        return jsonify({'error': '问题不能为空'})
    
    try:
        response = ask_deepseek(question, stream=True)
        if isinstance(response, str):  # 如果是错误信息
            return jsonify({'error': response})
        
        def generate():
            for line in response.iter_lines():
                if line:
                    try:
                        data = json.loads(line.decode('utf-8'))
                        if 'response' in data:
                            yield f"data: {json.dumps({'response': data['response']})}\n\n"
                        if data.get('done', False):
                            break
                    except json.JSONDecodeError:
                        continue
        
        return Response(generate(), mimetype='text/plain')
    except Exception as e:
        return jsonify({'error': f'流式响应错误：{str(e)}'})

if __name__ == '__main__':
    app.run(debug=True, host='0.0.0.0', port=5000)